#!/usr/bin/env python3
from scapy.all import *
import argparse
import time
import sys

def send_packet(target, protocol, ip_version, timeout=1):
    """
    发送单个数据包并等待确认
    """
    # 构造 IP 层
    if ip_version == 4:
        ip_layer = IP(dst=target)
    elif ip_version == 6:
        ip_layer = IPv6(dst=target)
    else:
        raise ValueError("Invalid IP version (must be 4 or 6)")

    # 构造传输层
    if protocol == "icmp":
        if ip_version == 4:
            packet = ip_layer/ICMP()
        else:
            # ip版本为6时使用如下的报文
            packet = ip_layer/ICMPv6EchoRequest()
        expected_response = ICMP
    elif protocol == "tcp":
        packet = ip_layer/TCP(dport=80, flags="S")  # SYN 包
        expected_response = TCP
    elif protocol == "udp":
        packet = ip_layer/UDP(dport=53)  # DNS 端口示例
        expected_response = None  # UDP 无确认
    else:
        raise ValueError("Unsupported protocol")

    # 发送并等待响应
    try:
        if expected_response:
            ans, unans = sr(packet, timeout=timeout, verbose=0)
            if ans:
                print(f"Received response from {target}")
                return True
            else:
                print(f"Timeout: No response from {target} in {timeout}s")
                return False
        else:
            send(packet, verbose=0)
            print(f"Sent {protocol.upper()} packet to {target} (no response expected)")
            return True
    except Exception as e:
        print(f"Error sending packet: {str(e)}")
        return False

def rate_controlled_send(target, protocol, ip_version, rate, duration):
    """
    速率控制发送
    """
    start_time = time.time()
    packets_sent = 0

    while (time.time() - start_time) < duration:
        batch_start = time.time()
        
        # 发送一个批次的包（每秒 rate 个）
        for _ in range(rate):
            send_packet(target, protocol, ip_version)
            packets_sent += 1

        # 等待至下一秒
        elapsed = time.time() - batch_start
        if elapsed < 1:
            time.sleep(1 - elapsed)

    print(f"\nTotal packets sent: {packets_sent}")

def main():
    parser = argparse.ArgumentParser(description="Scapy Packet Sender")
    parser.add_argument("target", help="Target IP address (IPv4/IPv6)")
    parser.add_argument("-p", "--protocol", choices=["icmp", "tcp", "udp"], 
                        default="icmp", help="Protocol to use")
    parser.add_argument("-6", "--ipv6", action="store_true",
                        help="Use IPv6 (default: IPv4)")
    parser.add_argument("-r", "--rate", type=int, default=1,
                        help="Packets per second (only for rate-controlled mode)")
    parser.add_argument("-d", "--duration", type=int, default=10,
                        help="Duration in seconds (only for rate-controlled mode)")
    parser.add_argument("-w", "--wait", action="store_true",
                        help="Wait for acknowledgment")

    args = parser.parse_args()

    # 确定 IP 版本
    ip_version = 6 if args.ipv6 else 4

    try:
        if args.wait:
            # 单次发送+等待确认模式
            success = send_packet(args.target, args.protocol, ip_version)
            sys.exit(0 if success else 1)
        else:
            # 速率控制模式
            rate_controlled_send(args.target, args.protocol, ip_version, 
                               args.rate, args.duration)
    except PermissionError:
        print("Error: Requires root privileges to send raw packets!")
        sys.exit(1)

if __name__ == "__main__":
    main()